import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from Res18 import ResNet18


def evaluate_model(model, test_loader, device):
    model.eval()
    criterion = nn.CrossEntropyLoss()

    total_loss = 0.0
    total_top1 = 0
    total_top5 = 0
    total_samples = 0

    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)
            total_loss += loss.item() * images.size(0)

            # Top-1
            _, preds = outputs.topk(5, 1, True, True)
            total_top1 += (preds[:, 0] == labels).sum().item()

            # Top-5
            total_top5 += sum([label in pred for label, pred in zip(labels, preds)])

            total_samples += images.size(0)

    avg_loss = total_loss / total_samples
    top1_acc = 100.0 * total_top1 / total_samples
    top5_acc = 100.0 * total_top5 / total_samples

    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6  # M为单位

    print(f"Test Loss: {avg_loss:.4f}")
    print(f"Top-1 Accuracy: {top1_acc:.2f}%")
    print(f"Top-5 Accuracy: {top5_acc:.2f}%")
    print(f"Total Parameters: {total_params:.2f} M")

def main():
    # 加载 CIFAR-100 测试集
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),
    ])
    test_dataset = datasets.CIFAR100(root='./data', train=False, download=True, transform=transform)
    test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=4)

    # 初始化模型
    device = torch.device("cpu")
    model = ResNet18(num_classes=100).to(device)

    # 加载训练好的权重（可选）
    model.load_state_dict(torch.load("resnet18_cifar100_15.pth", map_location=torch.device('cpu')))


    # 评估模型
    evaluate_model(model, test_loader, device)

if __name__ == "__main__":
    main()
